clc;clear all;close all;
global alpha betta depr sigmaC sigmaL phi govexp omega taumax N T ALPHA popweight
global kstst lstst cstst rstst wstst kstart kistart tauK0
global Delta1 Delta2 lamda kg lam kmax restrict

%run mainbench.m;

load('benchN2.mat')

kistart = kj;
kstart = sum(popweight.*kistart)+kg;
kmax=3;
ALPHA=0.4861;
lam=0.6;
lamda=lam;
Delta1 = 1;
Delta2 = 1;
restrict=1;

rstst = (betta^(-1)-1)/(1-taumax)+depr; %r from Euler at the steady state

xst0=[kstart 0.2];
xst = fsolve('findstst',xst0,options)
kstst=xst(1);
gammast=xst(2);


estst = (rstst/alpha)^(1/(1-alpha))*kstst; % express total effective labor from r=F_k
for jj=2:N
ljpart(jj-1)=popweight(jj)*phi(jj)*lamda(jj-1)^(-sigmaC/sigmaL)*(phi(jj)/phi(1))^(1/sigmaL);
end
lstst=estst/(popweight(1)*phi(1)+sum(ljpart)); %express l1 using the expression for total effective labor and l2=f(lamda,l1)
for jj=2:N
    L2a(jj-1) = lamda(jj-1)^(-sigmaC/sigmaL)*(phi(jj)/phi(1))^(1/sigmaL); %last terms in (11) and f_l
end

wstst = (1-alpha)*kstst^alpha*estst^(-alpha); %w=F_e
cstst = 1/(popweight(1)+sum(popweight(2:N).*lamda))*(kstst^alpha*estst^(1-alpha) - depr*kstst - govexp); %c from resource constraint

T=150; %we assume that the economy reaches the steady state after T periods

tauK0 = taumax;

%initial guesses 
for i=1:T;
    time1 = i/T;
    X(i) = time1*(kstst) + (1-time1)*kstart; %k
    X(T+i) = lstst; %l
    if (N==2)
    X(2*T+i) = 0.1; %gamma
    end  
end;
X(3*T+1:3*T+1+2*(N-1)) = [Delta1 Delta2 lamda];

ALPHall=[1 0.8 0.7 0.6 0.5 0.4861 0.4 0.3467 0.3 0.25 0.2 0.19 0.188 0.1879783 0.1879782 0.187 0.18 0.15 0.1 0.07 0.05 0.03 0.02 0.015 0.01]

for ii=1:length(ALPHall)       

ALPHA=ALPHall(ii)

tauK0 = taumax;
tolf0 = 1e-7;

options = optimset('MaxFunEvals',10000000,'MaxIter',200,'TolFun',1e-15,'TolX',1e-15,'Disp','iter');
[X,check] = fsolve('resid3v',X,options)

XX(ii,:)=X;
k = X(1:T);
l = X(T+1:2*T);
gamma = X(2*T+1:3*T);
Delta1 = X(3*T+1);
Delta2 = X(3*T+2:3*T+2+N-2);
lamda = X(3*T+3+N-2:3*T+3+2*(N-2));
Delta1*kistart(1)+sum(Delta2.*kistart(2:N))
xst0=[k(T) gamma(T)];
xst = fsolve('findstst',xst0,options)
kstst=xst(1);
gammast=xst(2);

estst = (rstst/alpha)^(1/(1-alpha))*kstst; % express total effective labor from r=F_k
for jj=2:N
ljpart(jj-1)=popweight(jj)*phi(jj)*lamda(jj-1)^(-sigmaC/sigmaL)*(phi(jj)/phi(1))^(1/sigmaL);
end
lstst=estst/(popweight(1)*phi(1)+sum(ljpart)); %express l1 using the expression for total effective labor and l2=f(lamda,l1)
for jj=2:N
    L2a(jj-1) = lamda(jj-1)^(-sigmaC/sigmaL)*(phi(jj)/phi(1))^(1/sigmaL); %last terms in (11) and f_l
end

wstst = (1-alpha)*kstst^alpha*estst^(-alpha); %w=F_e
cstst = 1/(popweight(1)+sum(popweight(2:N).*lamda))*(kstst^alpha*estst^(1-alpha) - depr*kstst - govexp); %c from resource constraint
%FOC for capital at steady state
Fkkst=alpha*(alpha-1)*kstst^(alpha-2)*estst^(1-alpha);
mustst=gammast*cstst^(-sigmaC)*Fkkst*(1-taumax)/(1-depr+rstst-1/betta); %mu from FOC for capital

klast = [kstart X(1:T-1)]; %k_{t-1}
e=[l' (l'*(lamda.^(-sigmaC/sigmaL).*(phi(2:N)/phi(1))'.^(1/sigmaL)))]*(popweight'.*phi);
e=e';

r = alpha*klast.^(alpha-1).*e.^(1-alpha); %r=F_k
w = (1-alpha)*klast.^(alpha).*e.^(-alpha); %w=F_e
rnext = [r(2:end) rstst];
lnext = [X(T+2:2*T) lstst];
c(T)=(betta*(1+(rstst-depr)*(1-taumax))*cstst.^(-sigmaC))^(-1/sigmaC);
for jj=T-1:-1:1
    c(jj)=(betta*(1+(r(jj+1)-depr)*(1-taumax))*c(jj+1).^(-sigmaC))^(-1/sigmaC);
end
cnext = [c(2:end) cstst];

%check bound on tauk
constrtest = c.^(-sigmaC)./cnext.^(-sigmaC) - betta*(1+ (rnext-depr)*(1-taumax));

cistst = [cstst lamda*cstst]; %consumption of groups in the long run
listst = [lstst, L2(lamda,lstst,phi(2:N))']; %hours worked of groups in the long run
for tt=1:T
ci(tt,:) = [c(tt) lamda*c(tt)]; %consumption of groups during the transition
li(tt,:) = [l(tt), L2(lamda,l(tt),phi(2:N))']; %hours worked of groups during the transition
end

%lifetime utilities
if (sigmaC==1)
    V = sum(betta.^[0:T-1]'*ones(1,N).*(log(ci)- omega*li.^(1+sigmaL)/(1+sigmaL)))...
    + betta^(T)/(1-betta)*(log(cistst) - omega*listst.^(1+sigmaL)/(1+sigmaL));
else
    V = sum(betta.^[0:T-1]'*ones(1,N).*((ci.^(1-sigmaC)-1)/(1-sigmaC) - omega*li.^(1+sigmaL)/(1+sigmaL)))...
    + betta^(T)/(1-betta)*((cistst.^(1-sigmaC)-1)/(1-sigmaC) - omega*listst.^(1+sigmaL)/(1+sigmaL));
end

if sum(V>=Vsq)==2
    PI(ii)=1
else
    PI(ii)=0
end

residall(ii)=sum(resid3v(X).^2)
Vall(ii,:)=V
Vsq

%tax rates
tauKt = 1 - ((c(1:end-1)./c(2:end)).^(-sigmaC)/betta-1)./(r(2:end)-depr); %Euler
tauKt = [tauK0 tauKt];
plot(tauKt)
tauKt;
lamda_all(ii)=lamda;
Delta1_all(ii)=Delta1;
Delta2_all(ii)=Delta2;
term_all(ii)=1+ALPHA*lamda^(1-sigmaC)+(Delta1+Delta2*lamda)*(1-sigmaC);
kappa=lamda^(-sigmaC/sigmaL)*(phi(2)/phi(1))^(1/sigmaL);
terml_all(ii)=1+ALPHA*kappa^(1+sigmaL)+(Delta1+phi(2)/phi(1)*kappa*Delta2)*(1+sigmaL);
cst_all(ii)=cstst;
gammast_all(ii)=gammast;
must_all(ii)=mustst;

%foc for labor
klam=lamda^(-sigmaC/sigmaL)*(phi(2)/phi(1))^(1/sigmaL);
lamderl(ii)=-lamda^(-sigmaC/sigmaL)*(phi(2)/phi(1))^((1+sigmaL)/sigmaL)*betta^T*(1+sigmaL)*lstst^(sigmaL)/(sum(betta.^[0:T-1].*c.^(-sigmaC).*c)+betta^T/(1-betta)*cstst^(-sigmaC)*cstst...
    -sigmaC/sigmaL*lamda^(-(sigmaC+sigmaL)/sigmaL)*(phi(2)/phi(1))^((1+sigmaL)/sigmaL)*(sum(betta.^[0:T-1].*(-omega*l.^(sigmaL).*l))+betta^T/(1-betta)*(-omega)*lstst^sigmaL*lstst));
%foc for consumption
lamderc(ii)=-lamda*betta^T*(1-sigmaC)*cstst^(-sigmaC)/(sum(betta.^[0:T-1].*c.^(-sigmaC).*c)+betta^T/(1-betta)*cstst^(-sigmaC)*cstst...
    -sigmaC/sigmaL*lamda^(-(sigmaC+sigmaL)/sigmaL)*(phi(2)/phi(1))^((1+sigmaL)/sigmaL)*(sum(betta.^[0:T-1].*(-omega*l.^(sigmaL).*l))+betta^T/(1-betta)*(-omega)*lstst^sigmaL*lstst));

tauLt = 1 - omega/phi(1)*l.^sigmaL./(c.^(-sigmaC).*w); %consumption-leisure FOC
tauLtall(ii,:) = tauLt;

save('solN2PO_taumax.mat')

(Vall(ii,:)>=Vsq)

end

kistart = flip(kj);
phi=flip(phi);
X=XX(1,:)

ALPHall2=[1 0.8 0.7 0.6 0.5 0.4 0.3 0.2 0.15 0.1 0.07 0.05 0.03 0.02 0.015 0.01 0.007 0.005 0.003 0.002 0.001 0]

for ii=1:length(ALPHall2)       

ALPHA=ALPHall2(ii)

tauK0 = taumax;
tolf0 = 1e-7;

options = optimset('MaxFunEvals',10000000,'MaxIter',200,'TolFun',1e-15,'TolX',1e-15,'Disp','iter');
[X,check] = fsolve('resid3v',X,options)

XX2(ii,:)=X;
k = X(1:T);
l = X(T+1:2*T);
gamma = X(2*T+1:3*T);
Delta1 = X(3*T+1);
Delta2 = X(3*T+2:3*T+2+N-2);
lamda = X(3*T+3+N-2:3*T+3+2*(N-2));
Delta1*kistart(1)+sum(Delta2.*kistart(2:N))
xst0=[k(T) gamma(T)];
xst = fsolve('findstst',xst0,options)
kstst=xst(1);
gammast=xst(2);

estst = (rstst/alpha)^(1/(1-alpha))*kstst; % express total effective labor from r=F_k
for jj=2:N
ljpart(jj-1)=popweight(jj)*phi(jj)*lamda(jj-1)^(-sigmaC/sigmaL)*(phi(jj)/phi(1))^(1/sigmaL);
end
lstst=estst/(popweight(1)*phi(1)+sum(ljpart)); %express l1 using the expression for total effective labor and l2=f(lamda,l1)
for jj=2:N
    L2a(jj-1) = lamda(jj-1)^(-sigmaC/sigmaL)*(phi(jj)/phi(1))^(1/sigmaL); %last terms in (11) and f_l
end

wstst = (1-alpha)*kstst^alpha*estst^(-alpha); %w=F_e
cstst = 1/(popweight(1)+sum(popweight(2:N).*lamda))*(kstst^alpha*estst^(1-alpha) - depr*kstst - govexp); %c from resource constraint
%FOC for capital at steady state
Fkkst=alpha*(alpha-1)*kstst^(alpha-2)*estst^(1-alpha);
mustst=gammast*cstst^(-sigmaC)*Fkkst*(1-taumax)/(1-depr+rstst-1/betta); %mu from FOC for capital

klast = [kstart X(1:T-1)]; %k_{t-1}
e=[l' (l'*(lamda.^(-sigmaC/sigmaL).*(phi(2:N)/phi(1))'.^(1/sigmaL)))]*(popweight'.*phi);
e=e';

r = alpha*klast.^(alpha-1).*e.^(1-alpha); %r=F_k
w = (1-alpha)*klast.^(alpha).*e.^(-alpha); %w=F_e
rnext = [r(2:end) rstst];
lnext = [X(T+2:2*T) lstst];
c(T)=(betta*(1+(rstst-depr)*(1-taumax))*cstst.^(-sigmaC))^(-1/sigmaC);
for jj=T-1:-1:1
    c(jj)=(betta*(1+(r(jj+1)-depr)*(1-taumax))*c(jj+1).^(-sigmaC))^(-1/sigmaC);
end
cnext = [c(2:end) cstst];

%check bound on tauk
constrtest = c.^(-sigmaC)./cnext.^(-sigmaC) - betta*(1+ (rnext-depr)*(1-taumax));

cistst = [cstst lamda*cstst]; %consumption of groups in the long run
listst = [lstst, L2(lamda,lstst,phi(2:N))']; %hours worked of groups in the long run
for tt=1:T
ci(tt,:) = [c(tt) lamda*c(tt)]; %consumption of groups during the transition
li(tt,:) = [l(tt), L2(lamda,l(tt),phi(2:N))']; %hours worked of groups during the transition
end

%lifetime utilities
if (sigmaC==1)
    V = sum(betta.^[0:T-1]'*ones(1,N).*(log(ci)- omega*li.^(1+sigmaL)/(1+sigmaL)))...
    + betta^(T)/(1-betta)*(log(cistst) - omega*listst.^(1+sigmaL)/(1+sigmaL));
else
    V = sum(betta.^[0:T-1]'*ones(1,N).*((ci.^(1-sigmaC)-1)/(1-sigmaC) - omega*li.^(1+sigmaL)/(1+sigmaL)))...
    + betta^(T)/(1-betta)*((cistst.^(1-sigmaC)-1)/(1-sigmaC) - omega*listst.^(1+sigmaL)/(1+sigmaL));
end

if sum(V>=Vsq)==2
    PI(ii)=1
else
    PI(ii)=0
end

residall2(ii)=sum(resid3v(X).^2)
Vall2(ii,:)=V
flip(Vsq)

%tax rates
tauKt = 1 - ((c(1:end-1)./c(2:end)).^(-sigmaC)/betta-1)./(r(2:end)-depr); %Euler
tauKt = [tauK0 tauKt];
plot(tauKt)
tauKt;
lamda_all2(ii)=lamda;
Delta1_all2(ii)=Delta1;
Delta2_all2(ii)=Delta2;
term_all2(ii)=1+ALPHA*lamda^(1-sigmaC)+(Delta1+Delta2*lamda)*(1-sigmaC);
kappa=lamda^(-sigmaC/sigmaL)*(phi(2)/phi(1))^(1/sigmaL);
terml_all2(ii)=1+ALPHA*kappa^(1+sigmaL)+(Delta1+phi(2)/phi(1)*kappa*Delta2)*(1+sigmaL);
cst_all2(ii)=cstst;
gammast_all2(ii)=gammast;
must_all2(ii)=mustst;

%foc for labor
klam=lamda^(-sigmaC/sigmaL)*(phi(2)/phi(1))^(1/sigmaL);
lamderl(ii)=-lamda^(-sigmaC/sigmaL)*(phi(2)/phi(1))^((1+sigmaL)/sigmaL)*betta^T*(1+sigmaL)*lstst^(sigmaL)/(sum(betta.^[0:T-1].*c.^(-sigmaC).*c)+betta^T/(1-betta)*cstst^(-sigmaC)*cstst...
    -sigmaC/sigmaL*lamda^(-(sigmaC+sigmaL)/sigmaL)*(phi(2)/phi(1))^((1+sigmaL)/sigmaL)*(sum(betta.^[0:T-1].*(-omega*l.^(sigmaL).*l))+betta^T/(1-betta)*(-omega)*lstst^sigmaL*lstst));
%foc for consumption
lamderc(ii)=-lamda*betta^T*(1-sigmaC)*cstst^(-sigmaC)/(sum(betta.^[0:T-1].*c.^(-sigmaC).*c)+betta^T/(1-betta)*cstst^(-sigmaC)*cstst...
    -sigmaC/sigmaL*lamda^(-(sigmaC+sigmaL)/sigmaL)*(phi(2)/phi(1))^((1+sigmaL)/sigmaL)*(sum(betta.^[0:T-1].*(-omega*l.^(sigmaL).*l))+betta^T/(1-betta)*(-omega)*lstst^sigmaL*lstst));

tauLt = 1 - omega/phi(1)*l.^sigmaL./(c.^(-sigmaC).*w); %consumption-leisure FOC
tauLtall2(ii,:) = tauLt;

save('solN2PO_taumax.mat')

(Vall2(ii,:)>=Vsq)

end

% clc;clear all;close all;
% load('solN2PO_taumax.mat')

Vcap=Vall(:,1);
Vwor=Vall(:,2);

term1=(1-sigmaC)*((1-betta)*Vcap+omega*listst_sq(1)^(1+sigmaL)/(1+sigmaL));
epscap=((term1+1).^(1/(1-sigmaC))/cistst_sq(1)-1)*100;
term2=(1-sigmaC)*((1-betta)*Vwor+omega*listst_sq(2)^(1+sigmaL)/(1+sigmaL));
epswor=((term2+1).^(1/(1-sigmaC))/cistst_sq(2)-1)*100;

Vcap=Vall2(:,2);
Vwor=Vall2(:,1);

term1=(1-sigmaC)*((1-betta)*Vcap+omega*listst_sq(1)^(1+sigmaL)/(1+sigmaL));
epscap2=((term1+1).^(1/(1-sigmaC))/cistst_sq(1)-1)*100;
term2=(1-sigmaC)*((1-betta)*Vwor+omega*listst_sq(2)^(1+sigmaL)/(1+sigmaL));
epswor2=((term2+1).^(1/(1-sigmaC))/cistst_sq(2)-1)*100;

epscap_taumax=[flip(epscap2)' epscap']
epswor_taumax=[flip(epswor2)' epswor']

save('solN2PO_taumax.mat')


